"""
@Time: 2020/12/6 下午 5:12
@Author: jinzhuan
@File: re_predictor_test.py
@Desc: 
"""
import torch
from cognlp import *
from cognlp.io.processor.rex import REXProcessor

torch.cuda.set_device(4)
device = torch.device('cuda')
processor = REXProcessor(path='../data/re/trexfull/data')
model = Bert4Re(len(processor.vocabulary))
predictor = RePredictor(model, model_path='../data/re/trexfull/model/ace2005-new/2020-12-21-15:12:38/checkpoint-15799/model.pt',
                        vocabulary=processor.vocabulary, device=device, device_ids=[5, 6])
print(predictor.predict(
    "Steven Jobs founded Apple.",
    [0, 2], [3, 4]))
